import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)

# 游戏设置
GRID_SIZE = 20  # 每个网格的像素大小
GRID_WIDTH = 12  # 游戏区域宽度（网格数）
GRID_HEIGHT = 12  # 游戏区域高度（网格数）
WIDTH = GRID_WIDTH * GRID_SIZE  # 游戏区域像素宽度
HEIGHT = GRID_HEIGHT * GRID_SIZE  # 游戏区域像素高度
SCORE_HEIGHT = 50  # 分数区域高度
TOTAL_HEIGHT = HEIGHT + SCORE_HEIGHT  # 总窗口高度
SNAKE_SPEED = 10  # 蛇的移动速度
SPEED_MIN = 5  # 最小速度
SPEED_MAX = 20  # 最大速度
WALL_PASS = False  # 是否允许穿墙

# 创建屏幕
screen = pygame.display.set_mode((WIDTH, TOTAL_HEIGHT))
pygame.display.set_caption('贪吃蛇')
clock = pygame.time.Clock()

# 蛇类
class Snake:
    def __init__(self):
        self.positions = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
        self.direction = (1, 0)  # 初始向右移动
        self.length = 1
        
    def get_head_position(self):
        return self.positions[0]
    
    def move(self):
        head_x, head_y = self.get_head_position()
        dir_x, dir_y = self.direction
        
        # 计算新位置
        new_x = head_x + dir_x
        new_y = head_y + dir_y
        
        # 检查是否撞墙
        if not WALL_PASS:
            if new_x < 0 or new_x >= GRID_WIDTH or new_y < 0 or new_y >= GRID_HEIGHT:
                return False  # 游戏结束
        else:
            # 穿墙模式
            new_x = new_x % GRID_WIDTH
            new_y = new_y % GRID_HEIGHT
        
        # 检查是否撞到自己
        if (new_x, new_y) in self.positions[:-1]:
            return False  # 游戏结束
        
        self.positions.insert(0, (new_x, new_y))
        if len(self.positions) > self.length:
            self.positions.pop()
        return True
    
    def grow(self):
        self.length += 1
    
    def change_direction(self, new_direction):
        # 防止直接反向移动
        if (new_direction[0] * -1, new_direction[1] * -1) != self.direction:
            self.direction = new_direction
    
    def draw(self, surface):
        for position in self.positions:
            rect = pygame.Rect(position[0] * GRID_SIZE, position[1] * GRID_SIZE, 
                              GRID_SIZE, GRID_SIZE)
            pygame.draw.rect(surface, GREEN, rect)
            pygame.draw.rect(surface, BLACK, rect, 1)

# 食物类
class Food:
    def __init__(self):
        self.position = (0, 0)
        self.randomize_position()
    
    def randomize_position(self, snake_positions=None):
        # 如果没有提供蛇的位置信息，则默认为空列表
        if snake_positions is None:
            snake_positions = []
        
        # 计算所有可用位置（不在蛇身体上的位置）
        available_positions = []
        for x in range(GRID_WIDTH):
            for y in range(GRID_HEIGHT):
                position = (x, y)
                if position not in snake_positions:
                    available_positions.append(position)
        
        # 如果没有可用位置（极端情况，蛇占满了整个地图），则不更新位置
        if not available_positions:
            return
        
        # 从可用位置中随机选择一个
        self.position = random.choice(available_positions)
    
    def draw(self, surface):
        rect = pygame.Rect(self.position[0] * GRID_SIZE, self.position[1] * GRID_SIZE, 
                          GRID_SIZE, GRID_SIZE)
        pygame.draw.rect(surface, RED, rect)
        pygame.draw.rect(surface, BLACK, rect, 1)

# 绘制网格线函数
def draw_grid(surface):
    for x in range(0, WIDTH, GRID_SIZE):
        pygame.draw.line(surface, BLACK, (x, 0), (x, HEIGHT), 1)
    for y in range(0, HEIGHT, GRID_SIZE):
        pygame.draw.line(surface, BLACK, (0, y), (WIDTH, y), 1)

# 主游戏函数
def main():
    # 添加用户设置选项
    global GRID_WIDTH, GRID_HEIGHT, WIDTH, HEIGHT, TOTAL_HEIGHT, WALL_PASS, SNAKE_SPEED
    
    # 这里可以添加用户设置界面，暂时使用默认值
    # 可以在这里添加代码让用户选择网格数和是否穿墙
    
    snake = Snake()
    food = Food()
    score = 0
    
    # 添加一个变量来跟踪当前待执行的指令
    pending_direction = None
    
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                # 处理方向键
                if event.key == pygame.K_UP:
                    pending_direction = (0, -1)
                elif event.key == pygame.K_DOWN:
                    pending_direction = (0, 1)
                elif event.key == pygame.K_LEFT:
                    pending_direction = (-1, 0)
                elif event.key == pygame.K_RIGHT:
                    pending_direction = (1, 0)
                # 添加切换穿墙模式的快捷键 - 使用单独的条件判断，不与方向键混在一起
                if event.key == pygame.K_w:
                    WALL_PASS = not WALL_PASS
                    print(f"穿墙模式: {'开启' if WALL_PASS else '关闭'}")
                # 添加调整速度的快捷键
                elif event.key == pygame.K_PLUS or event.key == pygame.K_KP_PLUS or event.key == pygame.K_EQUALS:
                    # 增加速度（最大为SPEED_MAX）
                    
                    SNAKE_SPEED = min(SNAKE_SPEED + 1, SPEED_MAX)
                    print(f"速度: {SNAKE_SPEED}")
                elif event.key == pygame.K_MINUS or event.key == pygame.K_KP_MINUS:
                    # 减少速度（最小为SPEED_MIN）
                    
                    SNAKE_SPEED = max(SNAKE_SPEED - 1, SPEED_MIN)
                    print(f"速度: {SNAKE_SPEED}")
        
        # 处理待执行的方向指令
        if pending_direction is not None:
            snake.change_direction(pending_direction)
            pending_direction = None
        
        # 移动蛇
        if not snake.move():
            # 游戏结束
            print(f"游戏结束! 得分: {score}")
            pygame.quit()
            sys.exit()
        
        # 检查是否吃到食物
        if snake.get_head_position() == food.position:
            snake.grow()
            food.randomize_position(snake.positions)
            score += 1
        
        # 绘制
        screen.fill(WHITE)
        
        # 绘制游戏区域
        pygame.draw.rect(screen, WHITE, (0, 0, WIDTH, HEIGHT))
        
        # 绘制网格线
        draw_grid(screen)
        
        # 绘制蛇和食物
        snake.draw(screen)
        food.draw(screen)
        
        # 绘制分数区域（在棋盘外）
        pygame.draw.rect(screen, (240, 240, 240), (0, HEIGHT, WIDTH, SCORE_HEIGHT))
        pygame.draw.line(screen, BLACK, (0, HEIGHT), (WIDTH, HEIGHT), 2)
        
        # 显示分数（移到棋盘外）
        font = pygame.font.SysFont('simhei', 24)
        score_text = font.render(f'得分: {score}', True, BLACK)
        screen.blit(score_text, (20, HEIGHT + 15))
        
        # 显示穿墙模式状态
        wall_text = font.render(f'穿墙模式: {"开启" if WALL_PASS else "关闭"} (按W切换)', True, BLACK)
        screen.blit(wall_text, (WIDTH - 300, HEIGHT + 15))
        
        # 显示速度状态
        speed_text = font.render(f'速度: {SNAKE_SPEED} (按+/-调整)', True, BLACK)
        screen.blit(speed_text, (WIDTH - 300, HEIGHT + 35))
        
        pygame.display.update()
        clock.tick(SNAKE_SPEED)

# AI游戏渲染函数
def render_ai_game(snake, food, score, episode, steps, ai_info="AI正在游戏中"):
    # 绘制
    screen.fill(WHITE)
    
    # 绘制游戏区域
    pygame.draw.rect(screen, WHITE, (0, 0, WIDTH, HEIGHT))
    
    # 绘制网格线
    draw_grid(screen)
    
    # 绘制蛇和食物
    snake.draw(screen)
    food.draw(screen)
    
    # 绘制分数区域（在棋盘外）
    pygame.draw.rect(screen, (240, 240, 240), (0, HEIGHT, WIDTH, SCORE_HEIGHT))
    pygame.draw.line(screen, BLACK, (0, HEIGHT), (WIDTH, HEIGHT), 2)
    
    # 显示分数
    font = pygame.font.SysFont('simhei', 24)
    score_text = font.render(f'得分: {score}', True, BLACK)
    screen.blit(score_text, (20, HEIGHT + 15))
    
    # 显示游戏信息
    episode_text = font.render(f'游戏: {episode}', True, BLACK)
    screen.blit(episode_text, (WIDTH - 300, HEIGHT + 15))
    
    # 显示步数
    steps_text = font.render(f'步数: {steps}', True, BLACK)
    screen.blit(steps_text, (WIDTH - 300, HEIGHT + 35))
    
    # 显示AI信息
    ai_text = font.render(ai_info, True, BLACK)
    screen.blit(ai_text, (20, HEIGHT + 35))
    
    pygame.display.update()
    
    # 处理退出事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

if __name__ == "__main__":
    main()